package main

import (
	"context"
	"flag"
	libLog "log"
	"net"
	"os"
	"sync"

	"github.com/libp2p/go-libp2p"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/multiformats/go-multiaddr"
)

var (
	log         = libLog.New(os.Stderr, "", libLog.Llongfile)
	forwardAddr string
	serverAddr  string
)

func init() {
	flag.StringVar(&forwardAddr, "f", "127.0.0.1:22", "-f ip:port")
	flag.StringVar(&serverAddr, "s", "", "-s /ip4/ip/tcp/port/p2p/key")
	flag.Parse()
}

func main() {
	if serverAddr == "" {
		log.Fatalln("Need a server address")
	}

	h, err := libp2p.New(context.Background())
	if err != nil {
		log.Panic(err)
	}
	defer h.Close()

	// get server addr info from serverAddr
	serverMAddr, err := multiaddr.NewMultiaddr(serverAddr)
	if err != nil {
		log.Panic(err)
	}
	serverPI, err := peer.AddrInfoFromP2pAddr(serverMAddr)
	if err != nil {
		log.Panic(err)
	}

	// connect and create stream
	err = h.Connect(context.Background(), *serverPI)
	if err != nil {
		log.Panic(err)
	}
	for {
		stream, err := h.NewStream(context.Background(), serverPI.ID, "pf.nealfree.cf")
		if err != nil {
			log.Println(err)
			continue
		}

		rxBuffer := make([]byte, 6)
		n, err := stream.Read(rxBuffer)
		if err != nil {
			log.Println(err)
			stream.Reset()
			continue
		}

		if string(rxBuffer[:n]) != "accept" {
			log.Println("wrong magic word:", string(rxBuffer[:n]))
			stream.Reset()
			continue
		}

		conn, err := net.Dial("tcp4", forwardAddr)
		if err != nil {
			log.Println(err)
			stream.Reset()
			continue
		}

		var wg sync.WaitGroup

		wg.Add(1)
		// rx
		go func() {
			defer wg.Done()
			rxBuffer := make([]byte, 1<<20)
			for {
				n, err := stream.Read(rxBuffer)
				if err != nil || n == 0 {
					log.Println(err)
					stream.Reset()
					break
				}

				_, err = conn.Write(rxBuffer[:n])
				if err != nil {
					log.Println(err)
					stream.Reset()
					break
				}
			}
		}()

		wg.Add(1)
		// tx
		go func() {
			defer wg.Done()
			txBuffer := make([]byte, 1<<20)
			for {
				n, err := conn.Read(txBuffer)
				if err != nil || n == 0 {
					// timeout issue
					log.Println(err)
					stream.Reset()
					break
				}

				_, err = stream.Write(txBuffer[:n])
				if err != nil {
					log.Println(err)
					stream.Reset()
					break
				}
			}
		}()
		wg.Wait()
		stream.Reset()
		log.Println("stream reset")
	}
}
